#include "crypto_utils.h"
#include <ctype.h>

void handle_openssl_errors(void) {
    ERR_print_errors_fp(stderr);
}

int generate_sm2_keypair(const char *private_key_file, const char *public_key_file) {
    EVP_PKEY_CTX *ctx = NULL;
    EVP_PKEY *pkey = NULL;
    FILE *priv_file = NULL, *pub_file = NULL;
    int ret = 0;

    printf("Generating SM2 key pair...\n");

    ctx = EVP_PKEY_CTX_new_id(EVP_PKEY_EC, NULL);
    if (!ctx) {
        printf("Failed to create context\n");
        return -1;
    }

    if (EVP_PKEY_keygen_init(ctx) <= 0) {
        printf("Failed to initialize key generation\n");
        goto cleanup;
    }

    if (EVP_PKEY_CTX_set_ec_paramgen_curve_nid(ctx, NID_sm2) <= 0) {
        printf("SM2 curve not available, using prime256v1\n");
        if (EVP_PKEY_CTX_set_ec_paramgen_curve_nid(ctx, NID_X9_62_prime256v1) <= 0) {
            printf("Failed to set EC curve\n");
            goto cleanup;
        }
    }

    if (EVP_PKEY_keygen(ctx, &pkey) <= 0) {
        printf("Failed to generate key pair\n");
        goto cleanup;
    }

    priv_file = fopen(private_key_file, "wb");
    if (!priv_file) {
        printf("Cannot open private key file for writing\n");
        goto cleanup;
    }
    if (PEM_write_PrivateKey(priv_file, pkey, NULL, NULL, 0, NULL, NULL) != 1) {
        printf("Failed to write private key\n");
        goto cleanup;
    }
    fclose(priv_file);
    priv_file = NULL;

    pub_file = fopen(public_key_file, "wb");
    if (!pub_file) {
        printf("Cannot open public key file for writing\n");
        goto cleanup;
    }
    if (PEM_write_PUBKEY(pub_file, pkey) != 1) {
        printf("Failed to write public key\n");
        goto cleanup;
    }
    fclose(pub_file);
    pub_file = NULL;

    printf("SM2 key pair generated successfully:\n");
    printf("  Private key: %s\n", private_key_file);
    printf("  Public key: %s\n", public_key_file);
    
    ret = 1;

cleanup:
    if (ctx) EVP_PKEY_CTX_free(ctx);
    if (pkey) EVP_PKEY_free(pkey);
    if (priv_file) fclose(priv_file);
    if (pub_file) fclose(pub_file);
    return ret;
}

int sm4_encrypt(const unsigned char *plaintext, size_t plaintext_len,
                const unsigned char *key, unsigned char **ciphertext, size_t *ciphertext_len) {
    EVP_CIPHER_CTX *ctx = NULL;
    unsigned char iv[SM4_BLOCK_SIZE] = {0};
    int outlen, tmplen;
    int ret = 0;

    *ciphertext = malloc(plaintext_len + SM4_BLOCK_SIZE);
    if (!*ciphertext) {
        printf("Memory allocation failed\n");
        return -1;
    }

    ctx = EVP_CIPHER_CTX_new();
    if (!ctx) {
        printf("Failed to create cipher context\n");
        goto cleanup;
    }

    if (EVP_EncryptInit_ex(ctx, EVP_sm4_cbc(), NULL, key, iv) != 1) {
        printf("Failed to initialize encryption\n");
        goto cleanup;
    }

    if (EVP_EncryptUpdate(ctx, *ciphertext, &outlen, plaintext, plaintext_len) != 1) {
        printf("Encryption update failed\n");
        goto cleanup;
    }

    if (EVP_EncryptFinal_ex(ctx, *ciphertext + outlen, &tmplen) != 1) {
        printf("Encryption final failed\n");
        goto cleanup;
    }

    *ciphertext_len = outlen + tmplen;
    printf("SM4 encryption successful: %zu bytes -> %zu bytes\n", plaintext_len, *ciphertext_len);
    
    ret = 1;

cleanup:
    if (ctx) EVP_CIPHER_CTX_free(ctx);
    if (ret != 1 && *ciphertext) {
        free(*ciphertext);
        *ciphertext = NULL;
    }
    return ret;
}

int sm4_decrypt(const unsigned char *ciphertext, size_t ciphertext_len,
                const unsigned char *key, unsigned char **plaintext, size_t *plaintext_len) {
    EVP_CIPHER_CTX *ctx = NULL;
    unsigned char iv[SM4_BLOCK_SIZE] = {0};
    int outlen, tmplen;
    int ret = 0;

    printf("SM4 Decrypt: ciphertext_len=%zu\n", ciphertext_len);

    *plaintext = malloc(ciphertext_len + SM4_BLOCK_SIZE);
    if (!*plaintext) {
        printf("Memory allocation failed for plaintext\n");
        return -1;
    }

    ctx = EVP_CIPHER_CTX_new();
    if (!ctx) {
        printf("Failed to create cipher context\n");
        goto cleanup;
    }

    if (EVP_DecryptInit_ex(ctx, EVP_sm4_cbc(), NULL, key, iv) != 1) {
        printf("Failed to initialize decryption\n");
        goto cleanup;
    }

    if (EVP_DecryptUpdate(ctx, *plaintext, &outlen, ciphertext, ciphertext_len) != 1) {
        printf("Decryption update failed\n");
        goto cleanup;
    }

    if (EVP_DecryptFinal_ex(ctx, *plaintext + outlen, &tmplen) != 1) {
        printf("Decryption final failed - padding or key issue\n");
        *plaintext_len = outlen;
        printf("Got partial decryption: %d bytes\n", outlen);
        
        if (outlen > 0) {
            printf("Partial data (hex): ");
            for (int i = 0; i < outlen && i < 32; i++) {
                printf("%02x ", (*plaintext)[i]);
            }
            printf("\n");
        }
        goto cleanup;
    }

    *plaintext_len = outlen + tmplen;
    printf("SM4 decryption successful: %zu bytes -> %zu bytes\n", ciphertext_len, *plaintext_len);
    
    ret = 1;

cleanup:
    if (ctx) EVP_CIPHER_CTX_free(ctx);
    if (ret != 1 && *plaintext) {
        free(*plaintext);
        *plaintext = NULL;
    }
    return ret;
}

int sm2_encrypt(EVP_PKEY *pub_key, const unsigned char *plaintext, size_t plaintext_len,
                unsigned char **ciphertext, size_t *ciphertext_len) {
    EVP_PKEY_CTX *ctx = NULL;
    int ret = 0;

    ctx = EVP_PKEY_CTX_new(pub_key, NULL);
    if (!ctx) {
        printf("Failed to create SM2 context\n");
        return -1;
    }

    if (EVP_PKEY_encrypt_init(ctx) <= 0) {
        printf("Failed to initialize SM2 encryption\n");
        goto cleanup;
    }

    if (EVP_PKEY_encrypt(ctx, NULL, ciphertext_len, plaintext, plaintext_len) <= 0) {
        printf("Failed to get ciphertext length\n");
        goto cleanup;
    }

    *ciphertext = malloc(*ciphertext_len);
    if (!*ciphertext) {
        printf("Memory allocation failed\n");
        goto cleanup;
    }

    if (EVP_PKEY_encrypt(ctx, *ciphertext, ciphertext_len, plaintext, plaintext_len) <= 0) {
        printf("SM2 encryption failed\n");
        goto cleanup;
    }

    ret = 1;

cleanup:
    if (ctx) EVP_PKEY_CTX_free(ctx);
    if (ret != 1 && *ciphertext) {
        free(*ciphertext);
        *ciphertext = NULL;
    }
    return ret;
}

int sm2_decrypt(EVP_PKEY *priv_key, const unsigned char *ciphertext, size_t ciphertext_len,
                unsigned char **plaintext, size_t *plaintext_len) {
    EVP_PKEY_CTX *ctx = NULL;
    int ret = 0;

    ctx = EVP_PKEY_CTX_new(priv_key, NULL);
    if (!ctx) {
        printf("Failed to create SM2 context\n");
        return -1;
    }

    if (EVP_PKEY_decrypt_init(ctx) <= 0) {
        printf("Failed to initialize SM2 decryption\n");
        goto cleanup;
    }

    if (EVP_PKEY_decrypt(ctx, NULL, plaintext_len, ciphertext, ciphertext_len) <= 0) {
        printf("Failed to get plaintext length\n");
        goto cleanup;
    }

    *plaintext = malloc(*plaintext_len);
    if (!*plaintext) {
        printf("Memory allocation failed\n");
        goto cleanup;
    }

    if (EVP_PKEY_decrypt(ctx, *plaintext, plaintext_len, ciphertext, ciphertext_len) <= 0) {
        printf("SM2 decryption failed\n");
        goto cleanup;
    }

    ret = 1;

cleanup:
    if (ctx) EVP_PKEY_CTX_free(ctx);
    if (ret != 1 && *plaintext) {
        free(*plaintext);
        *plaintext = NULL;
    }
    return ret;
}

int sm2_sign(EVP_PKEY *priv_key, const unsigned char *message, size_t message_len,
             unsigned char **signature, size_t *signature_len) {
    EVP_MD_CTX *md_ctx = NULL;
    int ret = 0;

    md_ctx = EVP_MD_CTX_new();
    if (!md_ctx) {
        printf("Failed to create signature context\n");
        return -1;
    }

    if (EVP_DigestSignInit(md_ctx, NULL, EVP_sm3(), NULL, priv_key) != 1) {
        printf("Failed to initialize signature\n");
        goto cleanup;
    }

    if (EVP_DigestSignUpdate(md_ctx, message, message_len) != 1) {
        printf("Failed to update signature\n");
        goto cleanup;
    }

    if (EVP_DigestSignFinal(md_ctx, NULL, signature_len) != 1) {
        printf("Failed to get signature length\n");
        goto cleanup;
    }

    *signature = malloc(*signature_len);
    if (!*signature) {
        printf("Memory allocation failed\n");
        goto cleanup;
    }

    if (EVP_DigestSignFinal(md_ctx, *signature, signature_len) != 1) {
        printf("Failed to create signature\n");
        goto cleanup;
    }

    ret = 1;

cleanup:
    if (md_ctx) EVP_MD_CTX_free(md_ctx);
    if (ret != 1 && *signature) {
        free(*signature);
        *signature = NULL;
    }
    return ret;
}

int sm2_verify(EVP_PKEY *pub_key, const unsigned char *message, size_t message_len,
               const unsigned char *signature, size_t signature_len) {
    EVP_MD_CTX *md_ctx = NULL;
    int ret = 0;

    md_ctx = EVP_MD_CTX_new();
    if (!md_ctx) {
        printf("Failed to create verification context\n");
        return -1;
    }

    if (EVP_DigestVerifyInit(md_ctx, NULL, EVP_sm3(), NULL, pub_key) != 1) {
        printf("Failed to initialize verification\n");
        goto cleanup;
    }

    if (EVP_DigestVerifyUpdate(md_ctx, message, message_len) != 1) {
        printf("Failed to update verification\n");
        goto cleanup;
    }

    if (EVP_DigestVerifyFinal(md_ctx, signature, signature_len) == 1) {
        ret = 1;
    } else {
        printf("Signature verification failed\n");
    }

cleanup:
    if (md_ctx) EVP_MD_CTX_free(md_ctx);
    return ret;
}

EVP_PKEY *load_public_key(const char *filename) {
    FILE *fp = fopen(filename, "rb");
    EVP_PKEY *pkey = NULL;
    
    if (!fp) {
        printf("Cannot open public key file: %s\n", filename);
        return NULL;
    }
    
    pkey = PEM_read_PUBKEY(fp, NULL, NULL, NULL);
    fclose(fp);
    
    if (!pkey) {
        printf("Failed to load public key from: %s\n", filename);
    }
    
    return pkey;
}

EVP_PKEY *load_private_key(const char *filename) {
    FILE *fp = fopen(filename, "rb");
    EVP_PKEY *pkey = NULL;
    
    if (!fp) {
        printf("Cannot open private key file: %s\n", filename);
        return NULL;
    }
    
    pkey = PEM_read_PrivateKey(fp, NULL, NULL, NULL);
    fclose(fp);
    
    if (!pkey) {
        printf("Failed to load private key from: %s\n", filename);
    }
    
    return pkey;
}

int create_digital_envelope(const unsigned char *plaintext, size_t plaintext_len,
                           EVP_PKEY *sender_priv, EVP_PKEY *receiver_pub,
                           DigitalEnvelope *envelope) {
    unsigned char sm4_key[SM4_KEY_SIZE];
    int ret = 0;

    if (RAND_bytes(sm4_key, SM4_KEY_SIZE) != 1) {
        printf("Failed to generate random SM4 key\n");
        return -1;
    }

    printf("Generated SM4 key: ");
    for (int i = 0; i < SM4_KEY_SIZE; i++) {
        printf("%02x", sm4_key[i]);
    }
    printf("\n");

    if (sm4_encrypt(plaintext, plaintext_len, sm4_key, &envelope->ciphertext, &envelope->ciphertext_len) != 1) {
        printf("SM4 encryption failed\n");
        goto cleanup;
    }

    if (sm2_encrypt(receiver_pub, sm4_key, SM4_KEY_SIZE, &envelope->encrypted_key, &envelope->encrypted_key_len) != 1) {
        printf("SM2 encryption failed\n");
        goto cleanup;
    }

    if (sm2_sign(sender_priv, envelope->ciphertext, envelope->ciphertext_len, &envelope->signature, &envelope->signature_len) != 1) {
        printf("SM2 signing failed\n");
        goto cleanup;
    }

    ret = 1;

cleanup:
    if (ret != 1) {
        free_digital_envelope(envelope);
    }
    return ret;
}

int extract_digital_envelope(DigitalEnvelope *envelope,
                            EVP_PKEY *receiver_priv, EVP_PKEY *sender_pub,
                            unsigned char **plaintext, size_t *plaintext_len) {
    unsigned char *decrypted_key = NULL;
    size_t key_len;
    int ret = 0;

    printf("\n=== Debug Information ===\n");
    printf("Ciphertext length: %zu\n", envelope->ciphertext_len);
    printf("Encrypted key length: %zu\n", envelope->encrypted_key_len);
    printf("Signature length: %zu\n", envelope->signature_len);

    printf("Verifying signature...\n");
    if (sm2_verify(sender_pub, envelope->ciphertext, envelope->ciphertext_len, envelope->signature, envelope->signature_len) != 1) {
        printf("Signature verification failed\n");
        return -1;
    }
    printf("✓ Signature verification successful!\n");

    printf("Decrypting SM4 key with SM2...\n");
    if (sm2_decrypt(receiver_priv, envelope->encrypted_key, envelope->encrypted_key_len, &decrypted_key, &key_len) != 1) {
        printf("SM2 decryption failed\n");
        return -1;
    }

    printf("Decrypted key length: %zu (expected: %d)\n", key_len, SM4_KEY_SIZE);
    printf("Decrypted SM4 key (hex): ");
    for (size_t i = 0; i < key_len && i < 32; i++) {
        printf("%02x", decrypted_key[i]);
    }
    printf("\n");

    if (key_len != SM4_KEY_SIZE) {
        printf("Warning: Decrypted key length mismatch, using first %d bytes\n", SM4_KEY_SIZE);
    }

    unsigned char sm4_key[SM4_KEY_SIZE];
    if (key_len >= SM4_KEY_SIZE) {
        memcpy(sm4_key, decrypted_key, SM4_KEY_SIZE);
    } else {
        memcpy(sm4_key, decrypted_key, key_len);
        memset(sm4_key + key_len, 0, SM4_KEY_SIZE - key_len);
    }
    
    printf("Final SM4 key (hex): ");
    for (int i = 0; i < SM4_KEY_SIZE; i++) {
        printf("%02x", sm4_key[i]);
    }
    printf("\n");

    printf("Attempting SM4 decryption...\n");
    printf("Ciphertext to decrypt (%zu bytes): ", envelope->ciphertext_len);
    for (size_t i = 0; i < envelope->ciphertext_len && i < 16; i++) {
        printf("%02x", envelope->ciphertext[i]);
    }
    printf("...\n");
    
    if (sm4_decrypt(envelope->ciphertext, envelope->ciphertext_len, sm4_key, plaintext, plaintext_len) != 1) {
        printf("SM4 decryption failed\n");
        
        if (*plaintext && *plaintext_len > 0) {
            printf("Partial decrypted data (%zu bytes): ", *plaintext_len);
            for (size_t i = 0; i < *plaintext_len && i < 32; i++) {
                if (isprint((*plaintext)[i])) {
                    printf("%c", (*plaintext)[i]);
                } else {
                    printf("\\x%02x", (*plaintext)[i]);
                }
            }
            printf("\n");
        }
        
        free(decrypted_key);
        return -1;
    }

    printf("✓ SM4 decryption successful!\n");
    printf("Decrypted plaintext length: %zu bytes\n", *plaintext_len);
    
    ret = 1;
    free(decrypted_key);
    return ret;
}

void free_digital_envelope(DigitalEnvelope *envelope) {
    if (envelope->ciphertext) free(envelope->ciphertext);
    if (envelope->encrypted_key) free(envelope->encrypted_key);
    if (envelope->signature) free(envelope->signature);
    envelope->ciphertext = NULL;
    envelope->encrypted_key = NULL;
    envelope->signature = NULL;
    envelope->ciphertext_len = 0;
    envelope->encrypted_key_len = 0;
    envelope->signature_len = 0;
}
